package com.kusch;

import com.baomidou.mybatisplus.generator.config.OutputFile;
import com.baomidou.mybatisplus.generator.config.po.TableInfo;
import com.baomidou.mybatisplus.generator.engine.FreemarkerTemplateEngine;
import com.sun.istack.internal.NotNull;
import org.apache.commons.lang3.StringUtils;

import java.io.File;
import java.util.Map;

/**
 * 代码生成器支持自定义[DTO\VO等]模版
 */
public final class EnhanceFreemarkerTemplateEngine extends FreemarkerTemplateEngine {

    @Override
    protected void outputCustomFile(@NotNull Map<String, String> customFile,
                                    @NotNull TableInfo tableInfo,
                                    @NotNull Map<String, Object> objectMap) {
        String entityName = tableInfo.getEntityName();
        //获取定义的 其他包的包名
        String otherPath = this.getPathInfo(OutputFile.other);
        customFile.forEach((key, value) -> {
            String fileName;
            //在配置中只写了 dto 的包，还需要VO的包
            //在这里判断一下，如果模板名称是VO结尾的，就放到VO包下去
            if (StringUtils.endsWith(key, "VO.java")) {
                String newPath = otherPath.replace("dto", "vo");
                fileName = String.format(newPath + File.separator + entityName + "%s", key);
            } else {
                fileName = String.format(otherPath + File.separator + entityName + "%s", key);
            }
            this.outputFile(new File(fileName), objectMap, value,
                    this.getConfigBuilder().getInjectionConfig().isFileOverride());
        });
    }
}
